package SSF.OS.NetFlow;

import SSF.OS.*;
import com.renesys.raceway.DML.*;
import SSF.OS.TCP.*;
import SSF.Net.*;
import SSF.Util.Streams.*;
import SSF.OS.BGP4.Comm.*;
import java.io.*;
import java.util.*;
import SSF.OS.Timer;

//Update:
//    --Feb 12, 2001. Add "ALL_PROTOCOLS" static variable to IpFlowCollector
//    --   check input NIC as well when identify flows.
//    --changed attribute "filter_type" to "protocol_type"
//    --In config(ip, configuration), add support to src_mask, dst_mask, 
// so that in IpNetFlow, pre-aggregation on IP address is supported.
//    --It won't be a protocol session any more.
//    --The packets with ACK flag won't be ignored anymore.
//    --the ConcurrentModificationException caused in moveConnection() is
// fixed.
//    --In init(), I added the part to write the header of netflow. This has
// to be changed when we upgrade the program to "having one netflow file per
// timeline"
//    --Wed.May 31, 2000, change this class to an extension of ProtocolSession.
// Filter for IP packet, collecting TCP flow daat

/**
 * Monitor that collects flows of record type: "SSF.OS.NetFlow"
 * Revised from IpFlowCollector, give "domain information" in flow records.
 * <P>In the future, it may be used to replace IpFlowCollector completely.
 */

public class IpFlowCollector implements ProtocolMonitor
{
    public static String STREAM_TYPE="SSF.OS.NetFlow";

    /** debug switch */
    protected boolean debug = false;

    /** the ip session it's related */
    protected IP ipSess = null;

    /** default value for protocol_type: all protocols */
    public static final byte ALL_PROTOCOLS = (byte)255;

    /** protocol interested in */
    public byte protocol_type = ALL_PROTOCOLS;

    /** table of open connections */
    protected IpFlowTable conOpenTable = null;

    /** init open table capacity*/
    protected int initCapacity = 101;

    /** List of finished connections */
    protected LinkedList conFinishList = null;

    /** time interval that the dump_timer will fire, default = 30 sec*/
    long dumpInterval = SSF.Net.Net.seconds(30.0);

    /** time interval to check "imcomplete" connections in conOpenTable*/
    long checkOpenTableInterval = SSF.Net.Net.seconds(20.0);

    /** longest time that a flow can be inactive */
    float maxInactiveTime = (float)10.0;

    /** longest flow lasting time , default = 10 mins */
    float maxFlowTime = (float)600.0;

    /** timer, periodically dump the conFinishList to a file*/
    protected Timer dumpTimer =  null;

    /** timer, check the open connection table when fired */
    protected Timer checkTableTimer = null;

    /** the stream to dump the data to */
    StreamInterface theStream = null;

    /** the datatype code (used when sending data to stream)*/
    protected int dataTypeCode;

    /** the host code (used when sending data to stream)*/
    protected int hostCode;

    /** the suffix source IP mask*/
    protected byte srcMask = (byte)0xff;

    /** the suffix desination IP mask */
    protected byte dstMask = (byte)0xff;

    /** the NIC class used to in receive function*/
    protected Class nicClass = null;

    /** used to check the type of the message, now just used for BGP*/
    protected Class protocolMsgTypeClass = null;

    //some constant
    protected final short LOCAL = -1;

    /**
     * default constructor, does nothing 
     */
    public IpFlowCollector() {};

    /**
     * config the collector
     * MUST be called before run the collector
     * <UL>
     * <LI> debug -- debug switch: true/false
     * <LI> protocol_type -- the protocol that is interested. By now it could be "tcp", "udp", "ospf", "bgp" or "app"(everything). Default is "tcp"
     * <LI> dump_interval -- the time interval that the flow data will be dumped to disk. Default is 30 seconds.
     * <LI> check_interval -- the time interval of checking the flow table for long-last flows and finished flows, preparing them for the next dump operation. Default is 20 seconds.
     * <LI> max_inactive_time -- if this flow hasn't receive any new packet for this period time, it's deemed as "finished". Default is 10 seconds.
     * <LI> max_flow_time -- the longest time of any flow, it will be cut if not finished yet. Default is 600 seconds.
     * <LI> init_capacity -- the initial capacity of the hash table used to store flow records. Default is 101, it'd better be a prime.
     * <LI> suffix_src_mask -- a byte that is used as suffix source IP mask.
     * <LI> suffix_dst_mask -- a byte that is used as suffix destination IP mask.
     * </UL>
     */
    public void config(ProtocolSession ownerProtocol, Configuration cfg) throws configException
    {
	//super.config(cfg);
	try {
	    ipSess = (IP)ownerProtocol;
	} catch (ClassCastException cce) {
	    throw new configException("IpFlowCollector should be installed on IP.");
	}
	//-- debug version or not?
	String vals = (String)cfg.findSingle("debug");
	if (vals != null) {
	    if (vals.equals("true"))
		debug = true;
	}

	//-- find the filter type
	String type_of_filter = (String)cfg.findSingle("protocol_type");
	if (type_of_filter == null)
	    System.err.println("Warning: no protocol type specified for IP monitor. Default is all protocols");
	else {
	    if (type_of_filter.equals("all"))
		protocol_type = ALL_PROTOCOLS;
	    else if (type_of_filter.equals("tcp"))
		protocol_type = Protocols.TCP_PRTL_NUM;
	    else if (type_of_filter.equals("udp"))
		protocol_type = Protocols.UDP_PRTL_NUM;
	    else if (type_of_filter.equals("ospf"))
		protocol_type = Protocols.OSPF_PRTL_NUM;
	    else if (type_of_filter.equals("bgp")) {
		protocol_type = Protocols.BGP_PRTL_NUM;
		try {
		    protocolMsgTypeClass = Class.forName("SSF.OS.BGP4.Comm.Message");
		} catch (ClassNotFoundException cnfe) {
		    System.err.println(cnfe.getMessage());
		    cnfe.printStackTrace();
		}
	    }
	    else if (type_of_filter.equals("app"))
		protocol_type = Protocols.APP_PRTL_NUM;
	    else if (type_of_filter.equals("test"))
		protocol_type = Protocols.TEST_PRTL_NUM;
	    else {
		System.err.println("Warning: unknown monitor specified for IP");
		System.err.println("by default it's set to tcp");
	    }
	}

	//-- specify dumpInterval
	vals = (String)cfg.findSingle("dump_interval");
	if (vals != null) {
	    //-- not the default value
	    dumpInterval = 
		SSF.Net.Net.seconds(Float.valueOf(vals).floatValue());
	}

	//-- specify the time interval to check the conOpenTable
	vals = (String)cfg.findSingle("check_interval");
	if (vals != null) {
	    //-- not default
	    checkOpenTableInterval = 
		SSF.Net.Net.seconds(Float.valueOf(vals).floatValue());
	}

	//-- specify maxInactive time
	vals = (String)cfg.findSingle("max_inactive_time");
	if (vals !=null) {
	    //-- not default
	    maxInactiveTime = Float.valueOf(vals).floatValue();
	}

	//-- specify the longest flowtime before it's removed from conOpenTable
	vals = (String)cfg.findSingle("max_flow_time");
	if (vals != null) {
	    //-- not default
	    maxFlowTime = Float.valueOf(vals).floatValue();
	}

	//-- sepcify the initCapacity of the ConOpenTable
	vals = (String)cfg.findSingle("init_capacity");
	if (vals != null) {
	    initCapacity = Integer.valueOf(vals).intValue();
	}

	//-- specify the suffix source IP mask
	vals = (String)cfg.findSingle("suffix_src_mask");
	if (vals != null) {
	    srcMask = Byte.valueOf(vals).byteValue();
	}

	//-- specify the suffix destination IP mask
	vals = (String)cfg.findSingle("suffix_dst_mask");
	if (vals != null) {
	    dstMask = Byte.valueOf(vals).byteValue();
	}

	//-- init the nicClass that is used in receive function.
	try {
	    nicClass = Class.forName("SSF.Net.NIC");
	}
	catch (Exception any) {
	    System.err.println(any.getMessage());
	    any.printStackTrace();
	    throw new configException("Error to init nicClass");
	}
    }

    /**
     * allocate memory for tables
     * init the "theStream" recorder
     * start the timers.
     */
    public void init()
    {
	ProbeSession theProbe = null;
	//-- allocate memory for the tables
	conOpenTable = new IpFlowTable(protocol_type, initCapacity);
	conFinishList = new LinkedList();
	
	try {
	    theProbe = 
		(ProbeSession)(ipSess.inGraph().SessionForName("probe"));
	}
	catch (ProtocolException pe) {
	    System.err.println("Error when init IpFlowCollector");
	    System.err.println(pe.getMessage());
	    System.exit(1);
	}

	// preconnect, and  spawn the thread that do wrapup work.
	theStream = theProbe.GetRecorder(new WrapupIpFlowCollector()); 

	hostCode = theProbe.getHostCode();  // uses the NHI address 

	dataTypeCode = theStream.getRecordTypeCode(STREAM_TYPE);
	
	//-- init & start timers
	startTimer();
    }

    /**
     * init and start the timers needed in this collector
     */
    private void startTimer()
    {
	//-- debug info
	//if (debug)
	//  System.err.println("Init flowCollector dumpTimer ...");      

	//-- create timers if needed
	if (theStream != null) {
	    dumpTimer = new Timer(inGraph(), dumpInterval) {
		    public void callback() {
			//-- debug info
			//if (debug)
			//System.err.println("dump ...");
		
			//-- dump
			dumpToStream();

			//if (debug)
			//System.err.println("dump done.");
		
			//-- reset timer
			set(dumpInterval);
		    }
		};
	}

	//-- debug info
	//if (debug)
	//System.err.println("Init flowCollector checkTableTimer ...");

	//-- create check timer
	checkTableTimer = new Timer(inGraph(), checkOpenTableInterval) {
	    public void callback() {

		//-- move long-last/inactive connections to conFinishList
		moveConnection();

		//-- reset timer
		set(checkOpenTableInterval);
	    }
	};

	dumpTimer.set(dumpInterval);
	checkTableTimer.set(checkOpenTableInterval);
    }

    /**
     * receive
     * filt the IpPacket, collect the statistical
     * values of the interested protocol
     * @param packet received ipPacket
     * @param fromSession the Protocol session from which this packet is received.
     * @param toSession the Protocol session to which this packet will be pushed to.
     */
    public void receive(ProtocolMessage packet, ProtocolSession fromSession, 
			ProtocolSession toSession)
    {
	IpHeader ipPacket = (IpHeader)packet;
	IpNetFlow flow = null;
	IpNetFlow flowList = null;
	Long key;

	//-- check whether it's the protocol interested in.
	//if (protocol_type != ipPacket.PROTOCOL_NO)
	//  return;         //-- not the protocol wanted, pass it.

	if (protocol_type != ALL_PROTOCOLS) {
	    if (ipPacket.PROTOCOL_NO == Protocols.TCP_PRTL_NUM) {
		
		if (protocol_type != Protocols.TCP_PRTL_NUM
		    && protocol_type != Protocols.BGP_PRTL_NUM)
		    return;
		
		// TCP or protocols based on TCP, check further
		switch (protocol_type) {
		case Protocols.TCP_PRTL_NUM :
		    break;
		case Protocols.BGP_PRTL_NUM:
		    //check whether it's BGP, if it is, go on.
		    Object payLoad = ipPacket.payload().payload();
		    if (payLoad == null) return;
		    if (!protocolMsgTypeClass.isInstance(payLoad))
			return;
		    break;
		case Protocols.APP_PRTL_NUM:
		    //further check needed to make sure it is interested.
		    break;
		default:
		    return;      //not the ones that are interested.
		}
	    }
	    //if it is "APP_PRTL_NUM", it will pass without further change. !!!!
	    //"APP_PRTL_NUM" is for Applications, that is, everthing.
	    else if (ipPacket.PROTOCOL_NO != Protocols.APP_PRTL_NUM) {
		//in current implementation
		//no other protocols are based on either UDP or sOSPF use
		if (ipPacket.PROTOCOL_NO != protocol_type)
		    return;
	    }
	}

	//-- check whether it's from NIC, if not return.
	short input= LOCAL;
	short output= LOCAL;

	if (nicClass.isInstance(fromSession))
	    input = (short)(((NIC)fromSession).ID);
	if (nicClass.isInstance(toSession))
	    output = (short)(((NIC)toSession).ID);

	float curTime = getSimTime();

	//-- compute the key for this ipPacket
	key = IpNetFlow.makeKey(ipPacket);

	//-- check whether the packet belongs to an existing flow
	flowList = (IpNetFlow)conOpenTable.get(key);

	if (flowList == null) {
	    //-- new flow, add to this entry
	    //System.err.println("No entry, add new Flow");
	    flow = new IpNetFlow(ipPacket, input, output, curTime,
				 ((Host)inGraph()).nhi, srcMask, dstMask);
	    conOpenTable.put(key, flow);
	}
	else {
	    //-- a list that potentially has the flow we want was returned
	    //-- check the elt  one by one, try to find matched flow
	    //-- ! Notice that it also check whether the input / output NIC 
	    //-- is the same.
	    //-- If not, it means it might split into a new flow,
	    //-- or, this is the point where two flows merge,
	    //-- so keep checking the followig elements until you find the one
	    //-- or you create a new one.
	    for (flow = flowList; flow!=null; flow=(IpNetFlow)(flow.next)) {
		if (flow.acceptPacket(ipPacket) 
		    && (output == flow.output)
		    && (input==flow.input) 
		    && (curTime - flow.Last <= maxInactiveTime))
		    break;
	    }

	    if (flow == null) {
		//-- no match, add a new flow to this list
		//System.err.println("add New Flow");
		flow=new IpNetFlow(ipPacket, input, output, curTime,
				   ((Host)inGraph()).nhi, srcMask, dstMask);
		flow.next = flowList;
		conOpenTable.put(key, flow);
	    }
	    else {                //-- there is a match, just update
		//for debug
		//System.err.println("Update");
		flow.update(packet, curTime);
	    }
	}

	//-- check whether this is a TCP FIN packet
	if (protocol_type == Protocols.TCP_PRTL_NUM) {	    
	    //-- check whether this is a FIN packet, 	
	    if (((TCP_Header)(ipPacket.payload())).flags 
		== TCP_Header.FIN) {
		//-- if yes, it's the end of this TCP flow
		//-- move it to the connectoin finish List
		//??? I can do slightly more efficient here by NOT calling
		//??? the removeFlow function... the flowList and the flow
		//??? are both known by now.
		conOpenTable.removeFlow(key, flow);
		//conFinishList.addHead(flow);
		conFinishList.addLast(flow);
	    }
	}
    }


    /**
     * getSimTime: get the simulation time
     * it's the same function used in tcpSessionMaster.java
     */
    public float getSimTime()
    {
	return ipSess.inGraph().now()/(float)SSF.Net.Net.seconds(1);
    }

    /**
     * inGraph
     * @ret return the protocolGraph the related IP session is in
     */
    public ProtocolGraph inGraph()
    {
	return ipSess.inGraph();
    }
    /**
     * moveConnection
     * move long-last/inactive connections to the conFinishList
     */
    private void moveConnection()
    {
	NetFlow flowList = null;
	NetFlow flow, prevFlow;
	Long key;

	float curTime = getSimTime();

	//-- check every entry
	for (Iterator iter=conOpenTable.values().iterator(); iter.hasNext();) {
	    try {
	    flowList = (NetFlow)iter.next();
	    }
	    catch (ConcurrentModificationException cme) {
		System.err.println("ConcurrentMOdificationException in IpFlowCollector.moveConnection -- iter.next()");
		System.exit(1);
	    }

	    key = flowList.makeKey();

	    //-- check this list except the first element

	    prevFlow = flowList;
	    for (flow=prevFlow.next; flow!=null; flow=(flow.next)) {
		if (isOldConnection(flow, curTime)) {
		    prevFlow.next = flow.next;    //remove it from list
		    flow.next = null;
		    //conFinishList.addHead(flow);
		    conFinishList.addLast(flow);

		    flow = prevFlow;     //so that the next field is correct 
		}
		else {                   //keep the prevFlow going
		    prevFlow = prevFlow.next;
		}
	    }

	    
	    if (isOldConnection(flowList, curTime)) {
		flow = flowList;
		//-- update the connection open table
		if (flowList.next == null)
		    iter.remove();
		else flowList = flowList.next;

		flow.next = null;
		conFinishList.addLast(flow);
	    }	   
	}
    }

    /**
     * dumpToStream
     * dump the data in the conFinishList to the StreamInterface
     */
    private void dumpToStream()
    {
	NetFlow flow;
	byte[] bytes;

	for (ListIterator iter=conFinishList.listIterator(0);
		 iter.hasNext();) {
	    flow = (NetFlow)iter.next();

	    if (debug)
		System.err.println(flow.toString());
	    bytes = flow.dumpToBytes();
	    theStream.send(dataTypeCode, hostCode,
			   getSimTime(), bytes, 0, bytes.length);
	}


	// clear the list
	conFinishList.clear();
    }
    /**
     * judge whether this connection is long-last/inactive
     * @param flow the IpNetFlow being checked
     * @param curTime current time
     */
    protected boolean isOldConnection(NetFlow nflow, float curTime)
    {
	IpNetFlow flow=null;
	try {
	    flow = (IpNetFlow)nflow;
	}
	catch (ClassCastException cce) {
	    System.err.println("IpNetFlow is expected. Please overload 'isOldConnection' method in IpFlowCollector.");
	    cce.printStackTrace();
	    System.exit(-1);
	}
	//-- whether it's inactive long enough
	if (curTime - flow.Last >= maxInactiveTime)
	    return true;

	//-- whether it lasts long enough
	if (curTime - flow.First >= maxFlowTime)
	    return true;

	return false;
    }

    /**
     * an inner class that does the wrap up work
     */
    class WrapupIpFlowCollector implements Runnable
    {
	public void run()
	{
	    if (debug)
		System.err.println("IpFlowCollector wrapup@"+((Host)(ipSess.inGraph())).nhi);

	    // dump all data left.
	    NetFlow flowList=null;
	    byte[] bytes;
	    float curTime = getSimTime();

	    for (Iterator iter=conOpenTable.values().iterator(); 
		 iter.hasNext();) {
		flowList = (NetFlow)iter.next();

		//dump all the flows in this list.
		while (flowList != null) {
		    bytes = flowList.dumpToBytes();
		    if (debug)
			System.out.println(flowList.toString());
		    theStream.send(dataTypeCode, hostCode,
				   curTime, bytes, 0, bytes.length);
		    flowList = flowList.next;
		}
	    }
	    
	    dumpToStream();
	    
	}
    }
}








